import tiktoken
import requests
import json
from dotenv import load_dotenv
import os
# 加载.env文件中的环境变量
load_dotenv()

def openai_count_tokens(text: str, model: str = "gpt-3.5-turbo") -> int:
    """
    计算文本的token数量
    
    Args:
        text: 要计算的文本
        model: 使用的模型名称，默认为gpt-3.5-turbo
        
    Returns:
        int: token数量
    """
    try:
        # 获取对应模型的编码器
        encoding = tiktoken.encoding_for_model(model)
        
        # 计算token数量
        tokens = encoding.encode(text)
        token_count = len(tokens)
        
        return token_count
    except Exception as e:
        print(f"计算token时发生错误: {e}")
        return 0


def glm_tokenize_text(content, model="glm-4-plus"):
    """
    调用智谱AI的tokenizer API
    
    参数:
    api_key (str): API密钥
    content (str): 需要处理的文本内容
    model (str): 使用的模型名称，默认为'glm-4-plus'
    
    返回:
    dict: API的响应结果
    """
    # 加载.env文件中的环境变量
    load_dotenv()

    # API端点
    url = os.getenv("GLM_TOKENIZER_URL")
    api_key = os.getenv("GLM_TOKEN")
    # 请求头
    headers = {
        "Authorization": api_key,
        "Content-Type": "application/json"
    }
    
    # 请求体
    payload = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": content
            }
        ]
    }
    
    try:
        # 发送POST请求
        response = requests.post(url, headers=headers, json=payload)
        
        # 确保请求成功
        response.raise_for_status()
        
        # 返回JSON响应
        response_json = response.json()
        token_count = response_json.get("usage", {}).get("prompt_tokens")
        return token_count
    
    except requests.exceptions.RequestException as e:
        print(f"请求失败: {e}")
        return None

# # 使用示例
# if __name__ == "__main__":
#     api_key = "您的apiKey"  # 替换为您的实际API密钥
#     content = "欢迎来到智谱AI开放平台"
    
#     result = tokenize_text(api_key, content)
#     if result:
#         print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    text = "今天天气怎么样？"
    token_count = openai_count_tokens(text)
    print(f"OPENAI的token数量为: {token_count}")
    GLM_token_count = glm_tokenize_text(text)
    print(f"GLM的token数量为: {GLM_token_count}")